"""
@author: wangkang
@contact: wangkang@autobio.com.cn
@file: parse_json.py
@time: 2020/1/22 16:05
@desc: 构造用户训练的数据
"""

import os
from collections import Counter


class TransferData:
    def __init__(self):
        cur = '/'.join(os.path.abspath(__file__).split('/')[:-1])
        self.label_dict = {
            '处理方式': 'WAY',
            '仪器部位': 'BODY',
            '故障反应': 'MANIFESTATION',
            '故障诊断': 'DIAGNOSE',
            '医院': 'HOSPITAL',
        }

        self.cate_dict = {
            'O': 0,
            'WAY-B': 1,
            'WAY-I': 2,
            'BODY-B': 3,
            'BODY-I': 4,
            'MANIFESTATION-B': 5,
            'MANIFESTATION-I': 6,
            'DIAGNOSE-B': 7,
            'DIAGNOSE-I': 8,
            'HOSPITAL-B': 9,
            'HOSPITAL-I': 10,
        }
        self.origin_path = os.path.join(cur, 'data_origin')
        self.train_filepath = os.path.join(cur, 'train.txt')
        return

    def transfer(self):
        f = open(self.train_filepath, 'w+')
        count = 0
        for root, dirs, files in os.walk(self.origin_path):
            for file in files:
                filepath = os.path.join(root, file)
                if 'original' not in filepath:
                    continue
                label_filepath = filepath.replace('.txtoriginal', '')
                print(filepath, '\t\t', label_filepath)
                content = open(filepath, encoding='UTF-8').read().strip()
                res_dict = {}
                for line in open(label_filepath, encoding='UTF-8'):
                    res = line.strip().split('	')
                    start = int(res[1])
                    end = int(res[2]) - 1
                    label = res[3]
                    label_id = self.label_dict.get(label)
                    for i in range(start, end + 1):
                        if i == start:
                            label_cate = label_id + '-B'
                        else:
                            label_cate = label_id + '-I'
                        res_dict[i] = label_cate

                for indx, char in enumerate(content):
                    char_label = res_dict.get(indx, 'O')
                    print(char, char_label)
                    f.write(char + '\t' + char_label + '\n')
        f.close()
        return


if __name__ == '__main__':
    handler = TransferData()
    train_datas = handler.transfer()
